#import brian_no_units
from brian import *
from brian.library.random_processes import *
from brian.library.synapses import  *
from constants import *

N = 1
N_RG = N
N_PF = N
N_IN = N 
N_INRG = N 

N_total = N_RG + N_PF + N_IN + N_INRG
print N_total




def iext(t):
    if t <= 500*ms:
        return 0*uA/cm**2
    elif t <= 5500*ms:
        return (t/ms-500) * 25/5000 *uA/cm**2
    elif t <= 15000*ms:
        return (25-(t/ms-5500) * 25/5000) *uA/cm**2
    else:
        return 0 *uA/cm**2

eqsMN=Equations('''
dh/dt= (hinf-h)/Tauh : 1
dn/dt= (ninf-n)/Taun : 1
dmnap/dt = (mnapinf-mnap)/Taumnap : 1
dhnap/dt = (hnapinf-hnap)/Tauhnap : 1
dm/dt = (minf-m)/Taum : 1
dvm/dt = 1/C*(iext(t)-INa-INap-IK-Ileak) : mV

Tauh = 30/(exp((vm/mV+50)/15)+exp(-(vm/mV+50)/16))*ms : ms
Taun = 7/(exp((vm/mV+40)/40)+exp(-(vm/mV+40)/50))*ms : ms
Tauhnap = 1200/(cosh(vm/mV + 59)/16)*ms :ms
minf = 1/(1+exp((vm-Vhm)/Sm)) : 1
hinf = 1/(1+exp((vm-Vhh)/Sh)) : 1
ninf = 1/(1+exp((vm-Vhn)/Sn)) : 1
mnapinf = 1/(1+exp((vm-Vhmnap)/Smnap)) : 1
hnapinf = 1/(1+exp((vm-Vhhnap)/Shnap)) : 1

INap = GNap*mnap*hnap*(vm-ENa) : mA*umetre**-2
INa = GNa*m**3*h*(vm-ENa) : mA*umetre**-2
IK = GK*n**4*(vm-EK) : mA*umetre**-2
Ileak = gleak*(vm-Eleak) : mA*umetre**-2

GNa :  msiemens*cm**-2
GNap :  msiemens*cm**-2
GK :  msiemens*cm**-2
gleak :  msiemens*cm**-2
Eleak : mV
''')

Group=NeuronGroup(N_RG+N_PF+N_IN+N_INRG,model=eqsMN,
threshold=EmpiricalThreshold(threshold=-40*mV),
implicit=True)

if N_RG != 0:
    SG_RG = Group.subgroup(N_RG)
    SG_RG.GNa = 30 * msiemens*cm**-2
    SG_RG.GNap = 0.25 * msiemens*cm**-2
    SG_RG.GK = 1 * msiemens*cm**-2
    SG_RG.gleak = 0.1 * msiemens*cm**-2
    SG_RG.Eleak = -64 * mV

if N_PF != 0:
    SG_PF = Group.subgroup(N_PF)
    SG_PF.GNa = 30 * msiemens*cm**-2
    SG_PF.GNap = 0.1 * msiemens*cm**-2
    SG_PF.GK = 1.2 * msiemens*cm**-2
    SG_PF.gleak = 0.1 * msiemens*cm**-2
    SG_PF.Eleak = -64 * mV
    
if N_IN != 0:
    SG_IN = Group.subgroup(N_IN)
    SG_IN.GNa = 120 * msiemens*cm**-2
    SG_IN.GNap = 0 * msiemens*cm**-2
    SG_IN.GK = 100 * msiemens*cm**-2
    SG_IN.gleak = 0.51 * msiemens*cm**-2
    SG_IN.Eleak = -64 * mV
    
if N_INRG != 0:
    SG_INRG = Group.subgroup(N_INRG)
    SG_INRG.GNa = 120 * msiemens*cm**-2
    SG_INRG.GNap = 0 * msiemens*cm**-2
    SG_INRG.GK = 100 * msiemens*cm**-2
    SG_INRG.gleak = 0.51 * msiemens*cm**-2
    SG_INRG.Eleak = -57.5 * mV
    
 
Group.vm = -70*mV
Group.mnap = 1/(1+exp((Group.vm-Vhmnap)/Smnap))
Group.hnap = 1/(1+exp((Group.vm-Vhhnap)/Shnap))
Group.m = 1/(1+exp((Group.vm-Vhm)/Sm))
Group.h = 1/(1+exp((Group.vm-Vhh)/Sh))
Group.n = 1/(1+exp((Group.vm-Vhn)/Sn))
        
    
tracevmINRG=StateMonitor(SG_INRG,'vm',record=[0])
tracevmRG=StateMonitor(SG_RG,'vm',record=[0])
tracevmPF=StateMonitor(SG_PF,'vm',record=[0])
tracevmIN=StateMonitor(SG_IN,'vm',record=[0])


spikes = SpikeMonitor(P)

DUR = 2*second

figure(1)
ion()
run(1*ms)
subplot(411)
tracevmRG.plot(refresh=100*ms,showlast=DUR)
subplot(412)
tracevmPF.plot(refresh=100*ms,showlast=DUR)
subplot(413)
tracevmIN.plot(refresh=100*ms,showlast=DUR)
subplot(414)
tracevmINRG.plot(refresh=100*ms,showlast=DUR)
run(DUR+1*ms)
ioff()
show()